If I remember correctly, I think it's up to our implementation to check if it is canceled. So I think we'll get stuck in the queue.take() 🤔

It doesn't seem like FutureUtils.cancel() will do an interrupt.

This is how we've handled that in the past: https://github.com/elastic/elasticsearch/blob/8.13/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java#L253

The shutdown() method puts a noop task on the queue to ensure that it wakes up.

Can you double check the tests and make sure we're covering this case (we call shutdown and then await termination)?

Adjusted with Add NoopTask to wake up queue on shutdown

Can you double check the tests and make sure we're covering this case (we call shutdown and then await termination)?

AFAIU we always check that when calling submitShutdownRequest, right?

How about we move this to a finally block to ensure it gets called.

Moved with Cleanup in finally block

I think there's a small potential for an edge case here (if we go the noop task route to do a shutdown for the queue.take()). If an exception occurs in startHandlingRateLimitedTasks(), it could cause the requestQueue to be drained which could mean that it'd never get the noop task.

I'd have to think of a good way to solve that. Maybe we split up the cleanup methods so that this one doesn't drain the requestQueue, instead the processRequestQueue() would call a different cleanup() that'd handle doing that 🤔

@jonathan-buttner something like Only reject requests of the respective execution path (rate-limited v…?

Do we need the cast?

I think it would be better to replace uses of RequestTask with RejectableTask in this class, since there's no reason that the interface can't be used throughout.

Remove unnecessary cast

If we're shutting down we need to reject the task we pulled off the requestQueue

Here's an example of doing that: https://github.com/elastic/elasticsearch/blob/8.13/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java#L192

Adjusted with Reject request on shutdown

I think if we move this to a finally block we probably don't need it here and we could probably remove the return.

If the task was accepted we'll want to check is shutdown one last time to ensure we notify this task that we're shutting down.

Here's an example: https://github.com/elastic/elasticsearch/blob/8.13/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java#L302

For consistency how about we make this static and accept a RateLimitSettings object. Then we can use it in executeEnqueuedTaskInternal which has a similar check. Technically executeEnqueuedTaskInternal should never receive a non-rate limited task, but it'd probably be good to check just in case.

Good point! Adjusted with Reuse rateLimitSettingsEnabled check

This code is duplicated in RateLimitingEndpointHandler, with the only difference being the format of the error message. Would it be possible to extract it to a static method that both places could call? Rather than needing to use the id field directly for the service grouping, it can be derived from the task: Integer.toString(task.getRequestManager().rateLimitGrouping().hashCode())

There are also quite a few other places where we reject tasks due to shutdown that could be extracted to a method call to reduce code duplication and ensure consistency for the error messages.

I think the id (inferenceEntityId being the name of the inference endpoint AFAIU) and service grouping are two different things, so IMO we should keep id in the error messages, so it's clear, which endpoint failed to execute a request. I think the hashCode of the rateLimitGrouping could become quite cryptic for a clear log/error message.

But nothing speaks against extracting the common part to a static method 👍

Extract rejection logic to common static method

Sorry for the confusion, the id field in RateLimitingEndpointHandler is the service grouping hash code. The field is not well named in that respect.

This should be service [%s] for inference

I think it actually is: error message formatting. AFAIUhashCode is simply executed on the inferenceEntityId leading to 3355 in this case.

What I mean is that you're using String.format() to construct the error message and passing requestManager.rateLimitGrouping().hashCode() as the second argument to it, but the string in the first argument doesn't have any placeholders for that value to be used. Right now the test is passing because the rate limit grouping is a constant from run to run, but if it ever changes, then the hard-coded hash code of 3355 will no longer be correct and the test will fail.

Ah got it now, adjusted with Use string placeholder in assertion

Is this change necessary? The test still passes if execute() is called instead of submitTaskToRateLimitedExecutionPath().

Also, while I know you didn't add it in this PR, this test's name is at odds with what it's actually asserting. The name implies that we don't expect an exception when polling the queue to cause the service to terminate, but the test explicitly terminates the service as part of throwing the exception, then asserts that it's terminated. Is this test actually testing anything other than that calling shutdown() causes the service to shut down?

Adjusted the test to use execute again.

Is this test actually testing anything other than that calling shutdown() causes the service to shut down?

Good point, I've adjusted the test to assert that the service still runs after task.execute(...) threw an exception. I've used task.execute(...) instead of queue.poll(), because poll() is also used internally in AdjustableCapacityBlockingQueue when calling requestQueue.take() in processRequestQueue, which does in fact terminate the service and made the test green for the wrong reason. Commit: Adjust test to check that a throwing task does not terminate the service. (I'll add a similar test for the non-rate-limited execution path)

Speaking of that: executeTaskImmediately inside processRequestQueue handles any Exceptions thrown by the processed request task without terminating the service. AFAIU take() shouldn't throw except when it's calling thread is interrupted, which we need to handle explicitly anyway. I've adjusted the error message in the general Exception handler with Adjust error message in general exception handler to reflect that an exception here is not coming from a task or a task rejection, but is potentially a more fundamental issue leading to service termination. Just wanted to double-check, if my reasoning is correct here and we want this behavior?

If I understand correctly, prior to this change, when execute() was called, we would retrieve (or create) a RateLimitingEndpointHandler appropriate to the task's rate limit settings, then add the task to the queue associated with that handler. This meant that there was one queue per RateLimitingEndpointHandler, each with a capacity defined by the RequestExecutorServiceSettings, and that we wouldn't reject a task unless the queue for the handler associated with it was full.

With the new change, all tasks are first submitted to the single requestQueue queue, which has the same capacity as each of the queues managed by the handlers, meaning that we can begin rejecting tasks even though none of the queues associated with the handlers are full, effectively reducing the total number of requests we can process in a given time.

Would it be better to instead only add tasks with no rate limit settings to the request queue, and call submitTaskToRateLimitedExecutionPath() in the execute() method for tasks with rate limit settings? That way, we're not adding rate limited tasks to one queue just to later remove them and add them to another queue.

That's a very valid point @DonalEvans! I'll adjust the implementation, thanks for flagging

Only add non-rate-limited tasks to fast-path request queue

This message might be better as "Failed to send request for inference id [%s] because the request executor service has been shutdown" to make it consistent with the error we report in execute() if we try to queue a task when we're shut down.

Adjust error message when request gets rejected

Not strictly related to this PR, but since you're making changes in this class, could you rename the id field on RateLimitingEndpointHandler to be something more descriptive, like rateLimitGroupingId?

Rename id in RateLimitingEndpointHandler to rateLimitGroupingId

This should be using Strings.format() with [3355] replaced with [%s] and the other argument being requestManager.rateLimitGrouping().hashCode().

Use Strings.format(...) in assertion

This suppression and unused variable can be avoided using the thenAnswer() method, which I think is a little cleaner:

when(mockExecutorService.submit(any(Runnable.class))).thenAnswer(i -> mock(Future.class));

Use thenAnswer instead of suppression

The rateLimitingEnabled check should be retained

if (isEmbeddingsIngestInput(inferenceInputs) || rateLimitingEnabled(requestManager.rateLimitSettings())) {

nit: When I first saw the comparison I thought this was a class name. How about we make this all caps or noopTask.